#include "conv_fwd.cuh"
#include <gtest/gtest.h>
#include <limits>
#include <spdlog/spdlog.h>
#include <torch/torch.h>

bool TestVcoreConvFwd(int n, int c, int h, int w, int k, int r, int s, int ph,
                      int pw, int sh, int sw) {
  spdlog::info("Create tensors...");
  auto torch_dtype = torch::kFloat;

  int p = (h + 2 * ph - r) / sh + 1;
  int q = (w + 2 * pw - s) / sw + 1;

  torch::Tensor nhwc_golden_a = torch::rand({n, h, w, c}, torch_dtype);
  torch::Tensor nhwc_golden_b = torch::rand({k, r, s, c}, torch_dtype);

  auto conv_option = torch::nn::functional::Conv2dFuncOptions()
                         .stride({sh, sw})
                         .padding({ph, pw})
                         .dilation({1, 1})
                         .groups(1);

  auto nchw_a = nhwc_golden_a.permute({0, 3, 1, 2});
  auto nchw_b = nhwc_golden_b.permute({0, 3, 1, 2});
  auto nchw_golden_c =
      torch::nn::functional::conv2d(nchw_a, nchw_b, conv_option);
  auto nhwc_golden_c = nchw_golden_c.permute({0, 2, 3, 1});
  auto nhwc_c = torch::zeros_like(nhwc_golden_c).cuda();

  auto nhwc_a = nhwc_golden_a.clone().cuda();
  auto nhwc_b = nhwc_golden_b.clone().cuda();

  dim3 block(16, 16); // 每个块中包含的线程数
  dim3 grid((q + block.x - 1) / block.x, (p + block.y - 1) / block.y, k);
  cudaStream_t stream;
  cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking);
  convolution2D_NHWC_kernel<float, float, float, float>
      <<<grid, block, 0, stream>>>(
          static_cast<const float *>(nhwc_a.data_ptr()),
          static_cast<const float *>(nhwc_b.data_ptr()),
          static_cast<float *>(nhwc_c.data_ptr()), n, h, w, c, k, r, s, p, q,
          sh, sw, ph, pw);
  cudaStreamSynchronize(stream);
  return true;
}

TEST(test_vcore_conv_fwd, EXAMPLE) {
  // EXPECT_TRUE(TestVcoreConvFwd(1, 3, 1, 1, 6, 1, 1, 0, 0, 1, 1));
  // EXPECT_TRUE(TestVcoreConvFwd(1, 2, 3, 3, 1, 3, 3, 1, 1, 2, 2));
  // EXPECT_TRUE(TestVcoreConvFwd(1, 3, 5, 5, 4, 3, 3, 1, 1, 2, 2));
  // EXPECT_TRUE(TestVcoreConvFwd(1, 3, 10, 10, 32, 3, 3, 1, 1, 2, 2));
  // EXPECT_TRUE(TestVcoreConvFwd(1, 3, 20, 20, 8, 3, 3, 1, 1, 2, 2));
  EXPECT_TRUE(TestVcoreConvFwd(1, 3, 380, 380, 48, 3, 3, 1, 1, 2, 2));
}